# volcano_plot.py
#
# Created: 2025-05-30
# Description: 绘制Science风格的火山图，带富集路径分层显示，用于分析差异表达基因。

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec


def plot_science_style_volcano(csv_path, output_path="volcano_plot.png"):
    """
    绘制带有富集通路条带的火山图，模拟Science风格。

    参数：
    - csv_path: str，CSV文件路径，要求包含列：'logFC'、'FDR'、'external_gene_name'
    - output_path: str，导出图片文件名（默认保存为 'volcano_plot.png'）
    """

    # ========== Step 1: 读取原始差异表达数据 ==========
    df = pd.read_csv(csv_path)  # 读取CSV数据为DataFrame格式

    # ========== Step 2: 数据预处理和派生字段构造 ==========

    # logFC是对数差异倍数，这里保留为log2FC列方便后续绘图
    df['log2FC'] = df['logFC']

    # 将FDR（校正后的p值）转换为-log10形式，以便表示显著性，避免对0取对数报错
    df['neg_log10_FDR'] = -np.log10(df['FDR'].replace(0, 1e-300))

    # 将基因名列统一为gene
    df['gene'] = df['external_gene_name']

    # 模拟pathway注释（前100为UP，100–200为DN，其余为None）
    df['pathway'] = 'None'
    df.loc[df.index[:100], 'pathway'] = 'KRAS_SIGNALING_UP'
    df.loc[df.index[100:200], 'pathway'] = 'KRAS_SIGNALING_DN'

    # 使用FDR排序得到全局显著性排名
    df['rank'] = df['FDR'].rank()

    # 构造分类：上下调（KRAS_UP/DOWN）+ 其他
    df['category'] = 'Other'
    df.loc[(df['log2FC'] < -0.5) & (df['rank'] <= 200), 'category'] = 'KRAS_UP'
    df.loc[(df['log2FC'] > 0.5) & (df['rank'] <= 200), 'category'] = 'KRAS_DOWN'

    # 给每个点分配颜色用于绘图（橙色为UP，蓝色为DN，灰色为Other）
    def assign_plot_color(row):
        if row['pathway'] == 'KRAS_SIGNALING_UP':
            return '#f5a623'  # 橙色表示KRAS上调
        elif row['pathway'] == 'KRAS_SIGNALING_DN':
            return '#4a90e2'  # 蓝色表示KRAS下调
        else:
            return '#999999'  # 灰色表示不显著或未注释

    df['plot_color'] = df.apply(assign_plot_color, axis=1)

    # 指定需要在图中特别标注的基因
    highlight_genes = ['KRAS', 'FOSL1', 'MYC']
    highlight_present = df[df['gene'].isin(highlight_genes)][['gene', 'log2FC', 'neg_log10_FDR']]

    # ========== Step 3: 图像布局设置（主图 + 条带图） ==========
    fig = plt.figure(figsize=(12, 10))  # 整体图尺寸
    gs = gridspec.GridSpec(4, 1, height_ratios=[10, 0.5, 0.5, 0.3], hspace=0.1)  # 分4个子图区域

    # 创建主图（火山图）和两个条带图（富集分布）
    ax_main = fig.add_subplot(gs[0])
    ax_strip_up = fig.add_subplot(gs[1], sharex=ax_main)
    ax_strip_dn = fig.add_subplot(gs[2], sharex=ax_main)
    ax_empty = fig.add_subplot(gs[3])  # 空白区用于布局对齐

    # 在主图中添加背景色区域，帮助区分显著上调与下调
    ax_main.axvspan(-2.3, -0.5, color='#DDE5FF', alpha=0.3)  # 蓝色背景：下调区域
    ax_main.axvspan(0.5, 2.3, color='#FFD8D5', alpha=0.3)  # 红色背景：上调区域

    # 分别筛选三类点（上调、下调、其他）
    df_other = df[df['plot_color'] == '#999999']
    df_kras_up = df[df['plot_color'] == '#f5a623']
    df_kras_dn = df[df['plot_color'] == '#4a90e2']

    # 主图绘制：灰色背景点
    ax_main.scatter(df_other['log2FC'], df_other['neg_log10_FDR'],
                    color='#999999', s=18, alpha=0.4)

    # 主图绘制：橙色显著上调点
    ax_main.scatter(df_kras_up['log2FC'], df_kras_up['neg_log10_FDR'],
                    color='#f5a623', s=25, edgecolor='gray', linewidth=0.2, alpha=0.9)

    # 主图绘制：蓝色显著下调点
    ax_main.scatter(df_kras_dn['log2FC'], df_kras_dn['neg_log10_FDR'],
                    color='#4a90e2', s=25, edgecolor='gray', linewidth=0.2, alpha=0.9)

    # 添加对称参考线
    ax_main.axvline(x=0, color='gray', linestyle='-', linewidth=1)

    # 添加Top 200显著区的虚线框（左侧UP，右侧DN）
    ax_main.add_patch(plt.Rectangle((-2.3, 2), 1.8, 5, linewidth=1,
                                    edgecolor='gray', facecolor='none', linestyle=':'))
    ax_main.add_patch(plt.Rectangle((0.5, 2), 1.8, 5, linewidth=1,
                                    edgecolor='gray', facecolor='none', linestyle=':'))

    # 添加区域标签
    ax_main.text(-1.5, 6.8, 'Top 200', fontsize=10, weight='bold')
    ax_main.text(1.3, 6.8, 'Top 200', fontsize=10, weight='bold')
    ax_main.text(-1.3, 7.2, 'PDAC KRAS UP\n(KRAS-dependent)', ha='center', fontsize=12, color='#4A90E2')
    ax_main.text(1.3, 7.2, 'PDAC KRAS DOWN\n(KRAS-inhibited)', ha='center', fontsize=12, color='#D3544B')

    # 添加显著性参考线（例如p=0.05）
    thresh = -np.log10(0.05)
    ax_main.axhline(y=thresh, color='gray', linestyle='--', linewidth=1)
    ax_main.text(3.1, thresh + 0.2, r'$\alpha$ = 0.05', fontsize=10)

    # 高亮特定基因名
    for _, row in highlight_present.iterrows():
        ax_main.text(row['log2FC'], row['neg_log10_FDR'], row['gene'],
                     fontsize=9, weight='bold',
                     bbox=dict(facecolor='white', edgecolor='gray', boxstyle='round,pad=0.2'))

    # 设置主图标题与Y轴标签
    ax_main.set_ylabel('Significance (-log10 adj. p-val.)')
    ax_main.set_title('Volcano Plot: Colored by Pathway (Science-style, layered)')

    # 设置主图Y轴范围
    ax_main.set_ylim(0, 8)

    # 不显示X轴刻度
    ax_main.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)

    # 设置每个子图边框为灰色
    for ax in [ax_main, ax_strip_up, ax_strip_dn]:
        for side in ['left', 'bottom', 'top', 'right']:
            ax.spines[side].set_color('gray')
            ax.spines[side].set_linewidth(1.2)

    # 条带图绘制：上调基因在log2FC上的位置（竖线）
    for x in df_kras_up['log2FC']:
        ax_strip_up.axvline(x=x, color='#f5a623', alpha=0.7, linewidth=0.8)

    # 条带图绘制：下调基因在log2FC上的位置（竖线）
    for x in df_kras_dn['log2FC']:
        ax_strip_dn.axvline(x=x, color='#4a90e2', alpha=0.7, linewidth=0.8)

    # 隐藏上方条带图的X轴
    ax_strip_up.tick_params(axis='x', which='both', bottom=False, top=False, labelbottom=False)

    # 设置下方条带图的X轴
    ax_strip_dn.set_xlabel('KRAS vs NS siRNA (log2FC)')
    ax_strip_dn.set_xlim(-3, 3)
    ax_strip_dn.set_xticks(np.arange(-3, 4, 1))
    ax_strip_dn.tick_params(axis='x', labelsize=10)

    # 隐藏条带图Y轴标签与刻度
    ax_strip_up.set_yticks([])
    ax_strip_dn.set_yticks([])
    ax_strip_up.set_ylabel('')
    ax_strip_dn.set_ylabel('')

    # 添加富集通路标签
    ax_strip_up.text(-2.95, 0.3, 'KRAS_SIGNALING_UP', fontsize=10, color='#f5a623', ha='left', va='center')
    ax_strip_dn.text(-2.95, 0.3, 'KRAS_SIGNALING_DN', fontsize=10, color='#4a90e2', ha='left', va='center')

    # 显示富集得分和p值
    ax_strip_up.text(2.7, 0.05, 'ES = -0.51\npval = 4.9e-12', fontsize=9, ha='right')
    ax_strip_dn.text(2.7, 0.05, 'ES = 0.54\npval = 1.6e-08', fontsize=9, ha='right')

    # 隐藏空白子图
    ax_empty.axis('off')

    # 调整子图布局，避免重叠
    plt.tight_layout()

    # 保存图像
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    plt.show()

if __name__ == '__main__':
    plot_science_style_volcano('Cancer_data_s1.csv', 'volcano_plot.png')